{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG with DocAgent\n",
    "\n",
    "In this notebook, see how DocAgent, through natural language, can\n",
    "\n",
    "1. Ingest documents from a local file or URL\n",
    "2. and answer questions with RAG capability\n",
    "\n",
    "For more information on DocAgent, read about [DocAgent in our Reference Agents](https://docs.ag2.ai/latest/docs/user-guide/reference-agents/docagent) section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation\n",
    "\n",
    "To get started with the document agent integration in AG2, follow these steps:\n",
    "\n",
    "Install AG2 with the `rag` extra:\n",
    "```bash\n",
    "pip install -U ag2[openai,rag]\n",
    "```\n",
    "\n",
    "Notes:\n",
    "\n",
    "   1. DocAgent only queries ingested documents, this ensures that it won't make up information if it can't find it.\n",
    "   2. Answers may not be accurate for documents that cannot be parsed correctly to Markdown format.\n",
    "\n",
    "You're all set! Now you can start using DocAgent feature in AG2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Documents supported\n",
    "\n",
    "The following documents can be ingested:\n",
    "- PDF\n",
    "- DOCX (DOCX, DOTX, DOCM, DOTM)\n",
    "- XLSX\n",
    "- PPTX (PPTX, POTX, PPSX, PPTM, POTM, PPSM)\n",
    "- HTML\n",
    "- ASCIIDOC (ADOC, ASCIIDOC, ASC)\n",
    "- MD (MD, MARKDOWN)\n",
    "- XML (XML, NXML)\n",
    "- TXT\n",
    "- JSON\n",
    "- CSV\n",
    "- IMAGE (BMP, JPG, JPEG, PNG, TIFF, TIF)\n",
    "\n",
    "You can also have the DocAgent use a web page by giving it a URL to ingest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inside the DocAgent\n",
    "\n",
    "![](https://media.githubusercontent.com/media/ag2ai/ag2/refs/heads/main/notebook/docagent_swarm.png)\n",
    "\n",
    "Before jumping into the code, let's have a look at the swarm inside the DocAgent.\n",
    "\n",
    "The swarm contains the following agents:\n",
    "\n",
    "- **Triage Agent**: Decides what type of task to perform from user requests.\n",
    "- **Task Manager Agent**: Manages the tasks and initiates actions.\n",
    "- **Data Ingestion Agent**: Ingests the documents.\n",
    "- **Query Agent**: Answers user questions based on ingested documents.\n",
    "- **Error Agent**: If anything fails, the error agent will report the problem back.\n",
    "- **Summary Agent**: Generates a summary of the completed tasks.\n",
    "\n",
    "When it is the DocAgent's turn to reply, the DocAgent initializes the agents and sets up the context variables, then:\n",
    "\n",
    "1. **Triage User Requests**: The `Triage Agent` categorizes the tasks into ingestions and queries.\n",
    "2. **Task Management**: The `Task Manager Agent` manages the tasks and ensures they are executed in the correct sequence.\n",
    "3. **Data Ingestion**: The `Data Ingestion Agent` processes any document ingesting tasks.\n",
    "4. **Query Execution**: The `Query Agent` answers any user queries.\n",
    "5. **Summary Generation**: The `Summary Agent` generates a summary of the completed tasks.\n",
    "\n",
    "The swarm is dynamic in that if only a data ingestation task is required, it will do that. Similarly for queries. This allows you to use the DocAgent many times covering the needs of ingesting, querying, or doing both at the same time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import autogen\n",
    "from autogen import AfterWorkOption, ConversableAgent, initiate_swarm_chat\n",
    "from autogen.agents.experimental import DocAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load configuration from the OAI_CONFIG_LIST file, requires the api_key to exist in the configuration\n",
    "config_list = autogen.config_list_from_json(\n",
    "    \"../OAI_CONFIG_LIST\",\n",
    "    filter_dict={\n",
    "        \"model\": [\"gpt-4.1\"],\n",
    "    },\n",
    ")\n",
    "\n",
    "# Set the OPENAI_API_KEY from the configuration so that the internal ingestion can use it\n",
    "os.environ[\"OPENAI_API_KEY\"] = config_list[0][\"api_key\"]\n",
    "\n",
    "llm_config = {\n",
    "    \"config_list\": config_list,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ingesting local documents and answering questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a document agent and ask them to ingest the document and answer the question\n",
    "document_agent = DocAgent(llm_config=llm_config, collection_name=\"toast_report\")\n",
    "run_response = document_agent.run(\n",
    "    message=\"could you ingest ../test/agentchat/contrib/graph_rag/Toast_financial_report.pdf? What is the fiscal year 2024 financial summary?\",\n",
    "    max_turns=1,\n",
    ")\n",
    "run_response.process()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [With Citations Support] Ingesting local documents and answering questions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen.agents.experimental.document_agent.chroma_query_engine import VectorChromaCitationQueryEngine\n",
    "\n",
    "# create a query engine with citation support\n",
    "query_engine = VectorChromaCitationQueryEngine(collection_name=\"toast_report\", enable_query_citations=True)\n",
    "# Create a document agent and ask them to ingest the document and answer the question\n",
    "document_agent = DocAgent(llm_config=llm_config, collection_name=\"toast_report\", query_engine=query_engine)\n",
    "run_response = document_agent.run(\n",
    "    message=\"could you ingest ../test/agentchat/contrib/graph_rag/Toast_financial_report.pdf? What is the fiscal year 2024 financial summary?\",\n",
    "    max_turns=1,\n",
    ")\n",
    "run_response.process()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fetching a webpage and answering questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a document agent and ask them to summarize a web page article\n",
    "document_agent = DocAgent(llm_config=llm_config, collection_name=\"news_reports\")\n",
    "run_response = document_agent.run(\n",
    "    message=\"could you read 'https://www.independent.co.uk/space/earth-core-inner-shape-change-b2695585.html' and summarize the article?\",\n",
    "    max_turns=1,\n",
    ")\n",
    "run_response.process()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple DocAgents in a Swarm\n",
    "\n",
    "Now we're going to use multiple DocAgents, each responsible for their own data.\n",
    "\n",
    "An nvidia_agent agent will ingest NVIDIA's financial report and query it. Similarly, an amd_agent will do the same with AMD's financial report.\n",
    "\n",
    "Although a single agent could ingest and query the documents, we want to ensure that their queries aren't tainted by the other company's documents. So we keep them separate and give them a unique `collection_name` which, in turn, will create individual data stores (See more on [Chroma collections](https://cookbook.chromadb.dev/core/collections/))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure the OPENAI_API_KEY is set in the environment\n",
    "llm_config = {\"model\": \"gpt-4.1\", \"api_type\": \"openai\", \"cache_seed\": None}\n",
    "\n",
    "# Create agents for the NVIDIA and AMD documents\n",
    "# Each agent has a unique collection_name so that data and queries are run in different vector store spaces\n",
    "nvidia_agent = DocAgent(\n",
    "    name=\"nvidia_agent\",\n",
    "    llm_config=llm_config,\n",
    "    collection_name=\"nvidia-demo\",\n",
    ")\n",
    "\n",
    "amd_agent = DocAgent(\n",
    "    name=\"amd_agent\",\n",
    "    llm_config=llm_config,\n",
    "    collection_name=\"amd-demo\",\n",
    ")\n",
    "\n",
    "# A financial analyst agent who will direct the DocAgents to ingest documents and answer questions\n",
    "# The financial analyst will also summarize the results and terminate the conversation\n",
    "analyst = ConversableAgent(\n",
    "    name=\"financial_analyst\",\n",
    "    system_message=(\n",
    "        \"You are a financial analyst working with two specialist agents, amd_agent who handles all AMD documents and queries, and nvidia_agent who handles all NVIDIA documents and queries. \"\n",
    "        \"Each agent knows how to load documents and answer questions from the document regarding their respective companies. \"\n",
    "        \"Only mention one of the two agents at a time, prioritize amd_agent. You will be able to engage each agent separately in subsequent iterations. \"\n",
    "        \"CRITICAL - Work with ONLY ONE agent at a time and provide (a) an indication for them to take action by saying '[Next to speak is ...]' together with (b) documents they need to ingest and (c) queries they need to run, if any. \"\n",
    "        \"DO NOT provide instructions that include the mention of both agents in the one response. \"\n",
    "        \"When all documents have been ingested and all queries have been answered, provide a summary and add 'TERMINATE' to the end of your summary. \"\n",
    "        \"The summary should contain detailed bullet points (multiple per query if needed) and grouped by each query. After the summary provide a one line conclusion. \"\n",
    "        \"Add the term '(This is not financial advice)' at the end of your conclusion. \"\n",
    "        \"If there are errors, list them and say 'TERMINATE'. \"\n",
    "        \"If there are no errors, do not say 'TERMINATE' until each agent has run their queries and provided their answers.\"\n",
    "    ),\n",
    "    is_termination_msg=lambda x: x.get(\"content\", \"\") and \"terminate\" in x.get(\"content\", \"\").lower(),\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "# Initiate the swarm (change the file paths in the messages if needed)\n",
    "result, _, _ = initiate_swarm_chat(\n",
    "    initial_agent=analyst,\n",
    "    agents=[analyst, nvidia_agent, amd_agent],\n",
    "    messages=(\n",
    "        \"Use the amd_agent to load AMD's 4th quarter 2024 report from \"\n",
    "        \"./docagent/AMDQ4-2024.pdf \"\n",
    "        \"and use the nvidia_agent to load NVIDIA's 3rd quarter 2025 report from \"\n",
    "        \"./docagent/NVIDIAQ3-2025.pdf. \"\n",
    "        \"Ask 'amd_agent' to ingest the AMD document and answer two queries (a) what AMD did in regards to AI and (b) what was the Q4 2024 GAAP revenue.\"\n",
    "        \"Ask 'nvidia_agent' to ingest the NVIDIA document and answer two queries (a) what NVIDIA did in regards to AI and (b) what was Q3 2025 GAAP revenue.\"\n",
    "    ),\n",
    "    swarm_manager_args={\n",
    "        \"llm_config\": llm_config,\n",
    "        \"system_message\": \"You are managing a financial analyst and two specialist company agents. After amd_agent or nvidia_agent, select the financial_analyst to speak next.\",\n",
    "        \"is_termination_msg\": lambda x: x.get(\"content\", \"\") and \"terminate\" in x.get(\"content\", \"\").lower(),\n",
    "    },\n",
    "    after_work=AfterWorkOption.SWARM_MANAGER,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tips for DocAgent\n",
    "\n",
    "- When asking for information from ingested information, be precise. For example, asking for revenue for the quarter in the previous example could retrieve a number of different revenue values, so we ask for \"GAAP revenue\" for the specific quarter.\n",
    "- If you have ingested documents in previous runs and just need to query the information, be clear to the DocAgent that they don't need to ingest documents you refer to.\n",
    "- Ensure that any files to be ingested can be accessed by the process you are running.\n",
    "- Ingestions take time, be sure to use `collection_name` to reuse collections that have already had the documents ingested in to.\n",
    "- You can review the Markdown files ingested in the `parsed_docs` folder to see how effective the conversion to Markdown was. This will help you investigate any query issues. "
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Query documents and web pages with DocAgent",
   "tags": [
    "agents",
    "documents",
    "RAG",
    "docagent"
   ]
  },
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
